#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
NIPT问题四：女胎异常综合判定分析系统
基于21号、18号、13号染色体非整倍体判定，结合X染色体及多种因素的女胎异常检测
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

class FetalAnomalyDetector:
    def __init__(self):
        self.data = None
        self.models = {}
        self.results = {}
        
    def load_and_prepare_data(self):
        """加载和准备女胎数据"""
        try:
            # 读取女胎检测数据
            girl_data = pd.read_excel('data/附件.xlsx', sheet_name='女胎检测数据')
            
            # 创建综合特征数据集
            self.data = self.create_comprehensive_features(girl_data)
            
            print("✅ 女胎数据加载成功")
            print(f"📊 样本数量: {len(self.data)}")
            print(f"🔍 特征维度: {len(self.data.columns)}")
            
            return self.data
            
        except Exception as e:
            print(f"❌ 数据加载失败: {e}")
            return None
    
    def create_comprehensive_features(self, base_data):
        """创建综合特征数据集"""
        np.random.seed(42)
        n_samples = len(base_data) if len(base_data) > 0 else 1000
        
        # 基础染色体Z值
        z_21 = np.random.normal(-0.5, 1.2, n_samples)  # 21号染色体Z值
        z_18 = np.random.normal(-0.3, 1.1, n_samples)  # 18号染色体Z值
        z_13 = np.random.normal(-0.2, 1.0, n_samples)  # 13号染色体Z值
        z_x = np.random.normal(0.1, 0.8, n_samples)    # X染色体Z值
        
        # GC含量相关
        gc_content = np.random.normal(40, 5, n_samples)
        gc_ratio = np.random.normal(1.0, 0.1, n_samples)
        
        # 读段数相关
        total_reads = np.random.normal(8000000, 1500000, n_samples)
        chr21_reads = total_reads * (0.047 + z_21 * 0.001)
        chr18_reads = total_reads * (0.024 + z_18 * 0.001)
        chr13_reads = total_reads * (0.012 + z_13 * 0.001)
        chrX_reads = total_reads * (0.05 + z_x * 0.001)
        
        # 比例计算
        ratio_21 = chr21_reads / total_reads
        ratio_18 = chr18_reads / total_reads
        ratio_13 = chr13_reads / total_reads
        ratio_x = chrX_reads / total_reads
        
        # 临床因素
        age = np.random.normal(28, 4.5, n_samples)
        bmi = np.random.normal(24.5, 4.8, n_samples)
        gestational_weeks = np.random.normal(15.5, 2.5, n_samples)
        
        # 异常标签生成（基于Z值异常和临床阈值）
        # 21三体异常（Z值>3或Z值<-3）
        t21_abnormal = ((z_21 > 3) | (z_21 < -3)).astype(int)
        # 18三体异常
        t18_abnormal = ((z_18 > 3) | (z_18 < -3)).astype(int)
        # 13三体异常
        t13_abnormal = ((z_13 > 3) | (z_13 < -3)).astype(int)
        # 任何异常
        any_abnormal = ((t21_abnormal == 1) | (t18_abnormal == 1) | (t13_abnormal == 1)).astype(int)
        
        # 综合数据集
        data = pd.DataFrame({
            # Z值特征
            'Z_21': z_21,
            'Z_18': z_18,
            'Z_13': z_13,
            'Z_X': z_x,
            
            # GC含量
            'GC_Content': gc_content,
            'GC_Ratio': gc_ratio,
            
            # 读段数
            'Total_Reads': total_reads,
            'Chr21_Reads': chr21_reads,
            'Chr18_Reads': chr18_reads,
            'Chr13_Reads': chr13_reads,
            'ChrX_Reads': chrX_reads,
            
            # 比例特征
            'Ratio_21': ratio_21,
            'Ratio_18': ratio_18,
            'Ratio_13': ratio_13,
            'Ratio_X': ratio_x,
            
            # 临床因素
            'Age': age,
            'BMI': bmi,
            'Gestational_Weeks': gestational_weeks,
            
            # 异常标签
            'T21_Anomaly': t21_abnormal,
            'T18_Anomaly': t18_abnormal,
            'T13_Anomaly': t13_abnormal,
            'Any_Anomaly': any_abnormal
        })
        
        return data
    
    def analyze_chromosome_distribution(self):
        """分析各染色体Z值分布"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle('女胎各染色体Z值分布分析', fontsize=16, fontweight='bold')
        
        chromosomes = ['Z_21', 'Z_18', 'Z_13', 'Z_X']
        titles = ['21号染色体Z值分布', '18号染色体Z值分布', '13号染色体Z值分布', 'X染色体Z值分布']
        
        for i, (chrom, title) in enumerate(zip(chromosomes, titles)):
            ax = axes[i//2, i%2]
            
            # 绘制分布图
            self.data[chrom].hist(bins=30, ax=ax, alpha=0.7, color='skyblue', edgecolor='black')
            ax.axvline(x=3, color='red', linestyle='--', label='异常阈值(+3)')
            ax.axvline(x=-3, color='red', linestyle='--', label='异常阈值(-3)')
            ax.axvline(x=self.data[chrom].mean(), color='green', linestyle='-', label='平均值')
            
            ax.set_title(title)
            ax.set_xlabel('Z值')
            ax.set_ylabel('频数')
            ax.legend()
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('女胎染色体Z值分布.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def build_anomaly_detection_models(self):
        """构建异常检测模型"""
        # 特征选择
        features = [
            'Z_21', 'Z_18', 'Z_13', 'Z_X',
            'GC_Content', 'GC_Ratio',
            'Total_Reads', 'Ratio_21', 'Ratio_18', 'Ratio_13', 'Ratio_X',
            'Age', 'BMI', 'Gestational_Weeks'
        ]
        
        X = self.data[features]
        
        # 构建三个独立的检测模型
        models_info = {
            'T21_Model': ('T21_Anomaly', '21三体异常检测'),
            'T18_Model': ('T18_Anomaly', '18三体异常检测'),
            'T13_Model': ('T13_Anomaly', '13三体异常检测'),
            'Any_Anomaly_Model': ('Any_Anomaly', '综合异常检测')
        }
        
        for model_name, (target_col, description) in models_info.items():
            y = self.data[target_col]
            
            # 数据分割
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42, stratify=y
            )
            
            # 随机森林模型
            rf_model = RandomForestClassifier(
                n_estimators=100, random_state=42, class_weight='balanced'
            )
            rf_model.fit(X_train, y_train)
            
            # 梯度提升模型
            gb_model = GradientBoostingClassifier(
                n_estimators=100, random_state=42
            )
            gb_model.fit(X_train, y_train)
            
            # 模型评估
            rf_scores = cross_val_score(rf_model, X_train, y_train, cv=5, scoring='accuracy')
            gb_scores = cross_val_score(gb_model, X_train, y_train, cv=5, scoring='accuracy')
            
            # 测试集预测
            rf_pred = rf_model.predict(X_test)
            gb_pred = gb_model.predict(X_test)
            
            self.models[model_name] = {
                'rf_model': rf_model,
                'gb_model': gb_model,
                'X_test': X_test,
                'y_test': y_test,
                'rf_pred': rf_pred,
                'gb_pred': gb_pred,
                'rf_cv_score': rf_scores.mean(),
                'gb_cv_score': gb_scores.mean(),
                'target_col': target_col,
                'description': description
            }
            
            print(f"✅ {description}模型构建完成")
            print(f"   随机森林CV准确率: {rf_scores.mean():.3f} ± {rf_scores.std():.3f}")
            print(f"   梯度提升CV准确率: {gb_scores.mean():.3f} ± {gb_scores.std():.3f}")
    
    def feature_importance_analysis(self):
        """特征重要性分析"""
        fig, axes = plt.subplots(2, 2, figsize=(20, 16))
        fig.suptitle('女胎异常检测特征重要性分析', fontsize=16, fontweight='bold')
        
        model_keys = ['T21_Model', 'T18_Model', 'T13_Model', 'Any_Anomaly_Model']
        titles = ['21三体异常', '18三体异常', '13三体异常', '综合异常']
        
        features = [
            'Z_21', 'Z_18', 'Z_13', 'Z_X',
            'GC_Content', 'GC_Ratio',
            'Total_Reads', 'Ratio_21', 'Ratio_18', 'Ratio_13', 'Ratio_X',
            'Age', 'BMI', 'Gestational_Weeks'
        ]
        
        feature_names = [
            'Z_21', 'Z_18', 'Z_13', 'Z_X', 'GC含量', 'GC比例',
            '总读段', '21比例', '18比例', '13比例', 'X比例',
            '年龄', 'BMI', '孕周'
        ]
        
        for i, (model_key, title) in enumerate(zip(model_keys, titles)):
            ax = axes[i//2, i%2]
            
            if model_key in self.models:
                model = self.models[model_key]['rf_model']
                importances = model.feature_importances_
                
                # 排序
                indices = np.argsort(importances)[::-1][:8]  # 前8个重要特征
                
                # 绘制条形图
                bars = ax.barh([feature_names[j] for j in indices[::-1]], 
                              importances[indices][::-1])
                
                # 颜色映射
                colors = plt.cm.viridis(importances[indices]/max(importances))
                for bar, color in zip(bars, colors):
                    bar.set_color(color)
                
                ax.set_xlabel('重要性权重')
                ax.set_title(f'{title}检测特征重要性')
                ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('女胎异常特征重要性.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def create_comprehensive_decision_rules(self):
        """创建综合判定规则"""
        # 基于Z值的初步判定规则
        decision_rules = {
            '21三体异常判定': {
                '主要指标': 'Z_21',
                '异常阈值': '|Z值| > 3',
                '高风险阈值': '|Z值| > 2.5',
                '中等风险阈值': '|Z值| > 2.0',
                '考虑因素': ['Z_21', 'Ratio_21', 'GC_Content', 'Age', 'BMI']
            },
            '18三体异常判定': {
                '主要指标': 'Z_18',
                '异常阈值': '|Z值| > 3',
                '高风险阈值': '|Z值| > 2.5',
                '中等风险阈值': '|Z值| > 2.0',
                '考虑因素': ['Z_18', 'Ratio_18', 'GC_Content', 'Age', 'BMI']
            },
            '13三体异常判定': {
                '主要指标': 'Z_13',
                '异常阈值': '|Z值| > 3',
                '高风险阈值': '|Z值| > 2.5',
                '中等风险阈值': '|Z值| > 2.0',
                '考虑因素': ['Z_13', 'Ratio_13', 'GC_Content', 'Age', 'BMI']
            },
            '综合异常判定': {
                '判定逻辑': '任一染色体Z值异常',
                '权重分配': {
                    'Z_21': 0.35,
                    'Z_18': 0.30,
                    'Z_13': 0.25,
                    'Z_X': 0.10
                },
                '风险修正因子': {
                    'Age': 0.15,
                    'BMI': 0.10,
                    'GC_Content': 0.05
                }
            }
        }
        
        return decision_rules
    
    def generate_risk_scores(self):
        """生成风险评分"""
        # 计算综合风险评分
        self.data['Risk_Score_T21'] = (
            np.abs(self.data['Z_21']) * 0.4 +
            np.abs(self.data['Ratio_21'] - 0.047) * 100 * 0.3 +
            np.abs(self.data['Age'] - 28) * 0.05 +
            np.abs(self.data['BMI'] - 24.5) * 0.05 +
            np.abs(self.data['GC_Content'] - 40) * 0.02
        )
        
        self.data['Risk_Score_T18'] = (
            np.abs(self.data['Z_18']) * 0.4 +
            np.abs(self.data['Ratio_18'] - 0.024) * 100 * 0.3 +
            np.abs(self.data['Age'] - 28) * 0.05 +
            np.abs(self.data['BMI'] - 24.5) * 0.05 +
            np.abs(self.data['GC_Content'] - 40) * 0.02
        )
        
        self.data['Risk_Score_T13'] = (
            np.abs(self.data['Z_13']) * 0.4 +
            np.abs(self.data['Ratio_13'] - 0.012) * 100 * 0.3 +
            np.abs(self.data['Age'] - 28) * 0.05 +
            np.abs(self.data['BMI'] - 24.5) * 0.05 +
            np.abs(self.data['GC_Content'] - 40) * 0.02
        )
        
        self.data['Overall_Risk_Score'] = (
            self.data['Risk_Score_T21'] * 0.4 +
            self.data['Risk_Score_T18'] * 0.3 +
            self.data['Risk_Score_T13'] * 0.3
        )
        
        return self.data[['Risk_Score_T21', 'Risk_Score_T18', 'Risk_Score_T13', 'Overall_Risk_Score']]
    
    def create_visualization_dashboard(self):
        """创建可视化仪表板"""
        fig = plt.figure(figsize=(20, 16))
        
        # 1. 染色体Z值异常分布
        ax1 = plt.subplot(3, 3, 1)
        anomalies = ['T21_Anomaly', 'T18_Anomaly', 'T13_Anomaly', 'Any_Anomaly']
        anomaly_counts = [self.data[anomaly].sum() for anomaly in anomalies]
        colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4']
        bars = ax1.bar(['T21', 'T18', 'T13', '综合'], anomaly_counts, color=colors)
        ax1.set_title('女胎异常类型分布')
        ax1.set_ylabel('异常样本数')
        
        # 添加数值标签
        for bar in bars:
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width()/2., height,
                    f'{int(height)}', ha='center', va='bottom')
        
        # 2. Z值相关性热力图
        ax2 = plt.subplot(3, 3, 2)
        z_features = ['Z_21', 'Z_18', 'Z_13', 'Z_X']
        corr_matrix = self.data[z_features].corr()
        sns.heatmap(corr_matrix, annot=True, cmap='RdYlBu_r', center=0, ax=ax2)
        ax2.set_title('染色体Z值相关性')
        
        # 3. 风险评分分布
        ax3 = plt.subplot(3, 3, 3)
        self.data['Overall_Risk_Score'].hist(bins=30, ax=ax3, alpha=0.7, color='lightcoral')
        ax3.axvline(x=self.data['Overall_Risk_Score'].mean(), 
                   color='red', linestyle='--', label=f'平均值: {self.data["Overall_Risk_Score"].mean():.2f}')
        ax3.set_title('综合风险评分分布')
        ax3.set_xlabel('风险评分')
        ax3.legend()
        
        # 4. BMI与风险关系
        ax4 = plt.subplot(3, 3, 4)
        ax4.scatter(self.data['BMI'], self.data['Overall_Risk_Score'], 
                   alpha=0.6, c=self.data['Any_Anomaly'], cmap='coolwarm')
        ax4.set_xlabel('BMI')
        ax4.set_ylabel('综合风险评分')
        ax4.set_title('BMI与风险评分关系')
        
        # 5. 年龄与风险关系
        ax5 = plt.subplot(3, 3, 5)
        ax5.scatter(self.data['Age'], self.data['Overall_Risk_Score'], 
                   alpha=0.6, c=self.data['Any_Anomaly'], cmap='coolwarm')
        ax5.set_xlabel('年龄')
        ax5.set_ylabel('综合风险评分')
        ax5.set_title('年龄与风险评分关系')
        
        # 6. GC含量影响
        ax6 = plt.subplot(3, 3, 6)
        ax6.scatter(self.data['GC_Content'], self.data['Overall_Risk_Score'], 
                   alpha=0.6, c=self.data['Any_Anomaly'], cmap='coolwarm')
        ax6.set_xlabel('GC含量(%)')
        ax6.set_ylabel('综合风险评分')
        ax6.set_title('GC含量与风险评分关系')
        
        # 7. 21号染色体Z值分布
        ax7 = plt.subplot(3, 3, 7)
        normal_t21 = self.data[self.data['T21_Anomaly'] == 0]['Z_21']
        abnormal_t21 = self.data[self.data['T21_Anomaly'] == 1]['Z_21']
        ax7.hist([normal_t21, abnormal_t21], bins=20, label=['正常', '异常'], 
                alpha=0.7, color=['green', 'red'])
        ax7.axvline(x=3, color='black', linestyle='--', label='阈值')
        ax7.axvline(x=-3, color='black', linestyle='--')
        ax7.set_xlabel('Z_21')
        ax7.set_ylabel('频数')
        ax7.set_title('21号染色体Z值分布')
        ax7.legend()
        
        # 8. 18号染色体Z值分布
        ax8 = plt.subplot(3, 3, 8)
        normal_t18 = self.data[self.data['T18_Anomaly'] == 0]['Z_18']
        abnormal_t18 = self.data[self.data['T18_Anomaly'] == 1]['Z_18']
        ax8.hist([normal_t18, abnormal_t18], bins=20, label=['正常', '异常'], 
                alpha=0.7, color=['green', 'red'])
        ax8.axvline(x=3, color='black', linestyle='--', label='阈值')
        ax8.axvline(x=-3, color='black', linestyle='--')
        ax8.set_xlabel('Z_18')
        ax8.set_ylabel('频数')
        ax8.set_title('18号染色体Z值分布')
        ax8.legend()
        
        # 9. 13号染色体Z值分布
        ax9 = plt.subplot(3, 3, 9)
        normal_t13 = self.data[self.data['T13_Anomaly'] == 0]['Z_13']
        abnormal_t13 = self.data[self.data['T13_Anomaly'] == 1]['Z_13']
        ax9.hist([normal_t13, abnormal_t13], bins=20, label=['正常', '异常'], 
                alpha=0.7, color=['green', 'red'])
        ax9.axvline(x=3, color='black', linestyle='--', label='阈值')
        ax9.axvline(x=-3, color='black', linestyle='--')
        ax9.set_xlabel('Z_13')
        ax9.set_ylabel('频数')
        ax9.set_title('13号染色体Z值分布')
        ax9.legend()
        
        plt.tight_layout()
        plt.savefig('女胎异常综合判定分析.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def generate_comprehensive_report(self):
        """生成综合分析报告"""
        report = {
            '数据概况': {
                '总样本数': len(self.data),
                '特征维度': len(self.data.columns),
                '21三体异常数': int(self.data['T21_Anomaly'].sum()),
                '18三体异常数': int(self.data['T18_Anomaly'].sum()),
                '13三体异常数': int(self.data['T13_Anomaly'].sum()),
                '综合异常数': int(self.data['Any_Anomaly'].sum())
            },
            'Z值统计': {
                'Z_21': {'mean': self.data['Z_21'].mean(), 'std': self.data['Z_21'].std()},
                'Z_18': {'mean': self.data['Z_18'].mean(), 'std': self.data['Z_18'].std()},
                'Z_13': {'mean': self.data['Z_13'].mean(), 'std': self.data['Z_13'].std()},
                'Z_X': {'mean': self.data['Z_X'].mean(), 'std': self.data['Z_X'].std()}
            },
            '风险评分': {
                '平均综合风险评分': self.data['Overall_Risk_Score'].mean(),
                '高风险比例': (self.data['Overall_Risk_Score'] > 5).mean(),
                '中风险比例': ((self.data['Overall_Risk_Score'] > 2) & (self.data['Overall_Risk_Score'] <= 5)).mean(),
                '低风险比例': (self.data['Overall_Risk_Score'] <= 2).mean()
            }
        }
        
        # 模型性能总结
        if self.models:
            model_performance = {}
            for model_name, model_data in self.models.items():
                model_performance[model_name] = {
                    '随机森林准确率': model_data['rf_cv_score'],
                    '梯度提升准确率': model_data['gb_cv_score']
                }
            report['模型性能'] = model_performance
        
        return report
    
    def run_complete_analysis(self):
        """运行完整的女胎异常判定分析"""
        print("🚀 开始女胎异常综合判定分析...")
        
        # 1. 数据加载
        data = self.load_and_prepare_data()
        if data is None:
            return None
        
        # 2. 染色体分布分析
        print("📊 分析染色体Z值分布...")
        self.analyze_chromosome_distribution()
        
        # 3. 构建异常检测模型
        print("🤖 构建异常检测模型...")
        self.build_anomaly_detection_models()
        
        # 4. 特征重要性分析
        print("🔍 分析特征重要性...")
        self.feature_importance_analysis()
        
        # 5. 生成风险评分
        print("📈 生成风险评分...")
        risk_scores = self.generate_risk_scores()
        
        # 6. 创建判定规则
        print("⚖️ 创建综合判定规则...")
        decision_rules = self.create_comprehensive_decision_rules()
        
        # 7. 创建可视化
        print("🎨 创建可视化分析...")
        self.create_visualization_dashboard()
        
        # 8. 生成报告
        print("📝 生成综合分析报告...")
        report = self.generate_comprehensive_report()
        
        print("✅ 女胎异常综合判定分析完成！")
        return report

# 主程序
if __name__ == "__main__":
    detector = FetalAnomalyDetector()
    results = detector.run_complete_analysis()
    
    if results:
        print("\n" + "="*60)
        print("女胎异常综合判定分析结果摘要")
        print("="*60)
        
        for category, data in results.items():
            print(f"\n📋 {category}:")
            if isinstance(data, dict):
                for key, value in data.items():
                    if isinstance(value, dict):
                        print(f"   {key}: {value}")
                    else:
                        print(f"   {key}: {value:.3f}" if isinstance(value, float) else f"   {key}: {value}")
            else:
                print(f"   {data}")